from math import log
import operator

def createDataSet():
    dataSet = [[1,1,'yes'],
                [1,1,'yes'],
                [1,0,'no'],
                [0,1,'no'],
                [0,1,'no']]
    labels = ['no surfacing','flippers']
    return dataSet,labels

def createTree(dataSet,labels):
    classList = [l[-1] for l in dataSet]
    if classList.count(classList[0]) == len(classList):
        return classList[0]
    if len(dataSet[0]) == 1:
        return majorityCnt(classList)
    bestFeat = chooseFeature(dataSet)
    bestLabel = labels[bestFeat]
    myTree = {bestLabel:{}}
    del(labels[bestFeat])
    feature = [l[bestFeat] for l in dataSet]
    unique = set(feature)
    for v in unique:
        subLabels = labels[:]
        myTree[bestLabel][v] = createTree(splitDataSet(dataSet,bestFeat,v),subLabels)
    return myTree

#计算香农信息熵
def calShannonEnt(dataSet):
    lens = len(dataSet)
    labelCount = {}
    for feat in dataSet:
        label = feat[-1]
        labelCount[label] = labelCount.get(label,0) + 1
    shannonEnt = 0.0
    for key in labelCount:
        prob = labelCount[key]/lens
        shannonEnt -= prob*log(prob,2)
    return shannonEnt

def splitDataSet(dataSet,axis,value):
    retDataSet = []
    for feature in dataSet:
        if feature[axis] == value:
            reducedFeat = feature[:axis]
            reducedFeat.extend(feature[axis+1:])
            retDataSet.append(reducedFeat)
    return retDataSet

def chooseFeature(dataSet):
    numFeats = len(dataSet[0]) - 1
    baseShanon = calShannonEnt(dataSet)
    baseGain = 0.0
    baseFeature = -1
    for i in range(numFeats):
        featlist = [l[i] for l in dataSet]#创建分类标签，计算标签的信息增益
        unique = set(featlist)
        newShanon = 0.0
        for v in unique:
            subData = splitDataSet(dataSet,i,v)
            prob = len(subData)/float(len(dataSet))
            newShanon += prob*calShannonEnt(subData)#计算信息熵
        infoGain = baseShanon - newShanon
        if(infoGain > baseGain):#计算最好的信息增益
            baseGain = infoGain
            baseFeature = i
    return baseFeature

def majorityCnt(classList):
    classCount = {}
    for vote in classList:
        classCount[vote] = classCount.get(vote,0) + 1
    sortedCount = sorted(classCount.items(),key = operator.itemgetter(1),reverse=True)
    return sortedCount[0][0]

def storeTree(inputTree,filename):#决策树的持久化
    import pickle
    f = open(filename,'w')
    pickle.dump(inputTree,f)
    f.close()

def loadTree(filename):
    import pickle
    f = open(filename,'r')
    return pickle.load(f)
if __name__ == '__main__':
    data,label = createDataSet()
    tree = createTree(data,label)
    print(tree)